from PyPDF2 import PdfReader # 读取pdf文件
from langchain.text_splitter import CharacterTextSplitter # 文本分割器
from langchain_community.vectorstores import FAISS # 向量库
from langchain_community.llms import QianfanLLMEndpoint # 千帆大模型平台库
import streamlit as st # 搭建web界面
from langchain.chains import ConversationalRetrievalChain # 对话检索链
from langchain.memory import ConversationBufferMemory
import os
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain.prompts import PromptTemplate

os.environ["QIANFAN_AK"] = '166LVKFb1FI2fg1St3rkJ2H9'
os.environ["QIANFAN_SK"] = 'anJgwBk8S4qcLfVzcKOhcSqn2edNNUm2'

st.title("《qa系统》")
st.write("请上传文档.")
# 设置上传pdf文件的功能
uploaded_file = st.file_uploader("选择一个pdf文档", type="pdf")

if uploaded_file:
# 读取pdf文件
    doc_reader = PdfReader(uploaded_file)
# 从pdf中提取文档
    raw_text = ""
    for i, page in enumerate(doc_reader.pages):
        text = page.extract_text()
        if text:
            raw_text += text
    #将文本切分成小的模块
    # print(raw_text)
    # print('*'*80)
    text_splitter = CharacterTextSplitter(separator="。", chunk_size=128, chunk_overlap=10)

    texts = text_splitter.split_text(raw_text)
    embeddings = QianfanEmbeddingsEndpoint()
    # 创建文档搜索
    docsearch = FAISS.from_texts(texts, embeddings)
    #提示
    prompt_template = """基于以下已知内容，简洁和专业的来回答用户的问题。
                                                    如果无法从中得到答案，清说"根据已知内容无法回答该问题"
                                                    答案请使用中文。
                                                    已知内容:
                                                    {context}
                                                    问题:
                                                    {question}"""
    prompt = PromptTemplate(template=prompt_template,
                            input_variables=["context", "question"])
    # 创建对话链
    qa = ConversationalRetrievalChain.from_llm(
                        llm=QianfanLLMEndpoint(model='Qianfan-Chinese-Llama-2-7B'),
                        retriever=docsearch.as_retriever(),
                        return_source_documents=True,
                        combine_docs_chain_kwargs={'prompt': prompt})
                        # 初始化聊天记录列表
    chat_history = []
                        # 获取用户的查询
    print(f'chat_history------------->{chat_history}')
    query = st.text_input("请给出你的问题")
    #添加一个生成按钮
    generate_button = st.button("生成答案")
    if generate_button and query:
        with st.spinner("答案生成中..."):
        # 将问题以及历史对话记录传入对话链获得模型输出结果
            result = qa({"question": query, "chat_history": chat_history})
            print(f'---------->{result}')
            answer = result["answer"]
            source_documents = result['source_documents']
            #将答案和source_documents合并为单个响应（输出）
            response = { "answer": answer, "source_documents": source_documents
            }
            st.write("response:", response)